import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np
import sys
sys.path.append("/home/project/Qtrainer_slim/test/yolov4") 
from configs import yolov4_cfg
from collections import OrderedDict
import torch.nn.functional as F

class YoloV4CspNet(nn.Module):

  
    def __init__(self):
        super(YoloV4CspNet,self).__init__()

        backboneConfig = yolov4_cfg.yolov4_cfg.configs
        # print("ddd",len(backboneConfig))
        
        self.layerList = locals()
        self.layerNameList = []
        # self.yolov4backbone = nn.Sequential()
        self.yolov4Layer = nn.Sequential()

        for i in range(len(backboneConfig)):
            index = backboneConfig[i][0]
            in_channels= backboneConfig[i][1]
            out_channels= backboneConfig[i][2]
            kernel_size= backboneConfig[i][3]
            stride= backboneConfig[i][4]
            padding= backboneConfig[i][5]
            momentum= backboneConfig[i][6]
            eps= backboneConfig[i][7]
            # “conv”,"route","shortcut","upsample","yolohead","maxpool" 
            layertype = backboneConfig[i][8]
            layersp = backboneConfig[i][9]

            layername = "layer"+"_"+str(index)
            self.layerNameList.append(layername)
            # 卷积块
            if(layertype == "conv"):
                # self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps))
                # self.layerList[layername] = Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)
                exec("self.layer"+"_"+str(index)+"=Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)")

                
                # self. = Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)
            # concate                 
            # elif(layertype == "route"):
            #     print("shortcut",layersp)
            #     routeLayer = {}
            #     for i in range(len(layersp)):
            #         routeName = "route_"+str(layersp[i])
            #         routeLayer[routeName] = self.yolov4Layer[layersp[i]]
            #         print("routeLayer", routeLayer)

            #     self.yolov4Layer.add_module(layername,torch.cat(**routeLayer))
            # # shortcut    
            # elif(layertype == "shortcut"):
            #     print("shortcut",layersp)
               
            #     self.yolov4Layer.add_module(layername,torch.add(self.yolov4Layer[layersp[0]],self.yolov4Layer[layersp[1]]))
            #         # torch.cat()

            #     # self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps))
            # # upsample    
            # elif(layertype == "upsample"):
            #     self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)) 
            # # yolohead    
            # elif(layertype == "yolohead"):
            #     self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps))
            # # maxpool     
            # elif(layertype == "maxpool"):
            #     self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)) 

            print("index",index)
            layername = index
            
            # self.layerList.append(layername)

        # self.fc1 = nn.Linear(32*40*40,2)
        self.upsampleLayer = torch.nn.Upsample(scale_factor=2, mode='bicubic', align_corners=True)

        self.maxPool = torch.nn.AdaptiveMaxPool1d(1)

        self.headlinerMax = nn.Linear(76*76*255,1)
        self.headlinerMedian = nn.Linear(38*38*255,1)
        self.headlinerSmall = nn.Linear(19*19*255,1)

    def forward(self,x):

        backboneConfig = yolov4_cfg.yolov4_cfg.configs
      
        result = []
        # 每层输出值
        layerValueList = []
        
        y = x
        print("self.layer12",self.layer_0)
        # print("self.layerList",self.layerList)
        # print("ddd",len(backboneConfig))
        for i in range(len(backboneConfig)):
            index = backboneConfig[i][0]
            in_channels= backboneConfig[i][1]
            out_channels= backboneConfig[i][2]
            kernel_size= backboneConfig[i][3]
            stride= backboneConfig[i][4]
            padding= backboneConfig[i][5]
            momentum= backboneConfig[i][6]
            eps= backboneConfig[i][7]
            # “conv”,"route","shortcut","upsample","yolohead","maxpool" 
            layertype = backboneConfig[i][8]
            layersp = backboneConfig[i][9]
            # layername = str(index)
            layername = "layer"+"_"+str(index)
           
            # 卷积块
            if(layertype == "conv"):
                # x = self.yolov4Layer[layername](x)
                # z = y
                # yi = torch.tensor([])
                # y = {}
                # yi = {"z": y}
                
                # exec("self.layer"+"_"+str(index)+"(z)",yi)
                # exec("y = self.layer"+"_"+str(index)+"(y)")
                y = eval("self.layer"+"_"+str(index)+"(y)")
                # print("yi",yi)
                # eval(self.layerNameList[index])(x)
                # x = getattr(a, "run")(x) 
                # print("conv_x",x)
                print("layername_conv",layername)
                print("y",y.shape)
                layerValueList.append(y)
            # concate                 
            elif(layertype == "route"):
                # print("shortcut",layersp)
                routeValueList = []
                for i in range(len(layersp)):
                    routeName = "route_"+str(layersp[i])
                    # print("routeName",routeName)
                    # print("lvl",layerValueList)
                    print("layerValueList["+str(layersp[i])+"]",layerValueList[layersp[i]].shape)
                    routeValueList.append(layerValueList[layersp[i]])
                    # routeValueList = routeValueList + layerValueList[layersp[i]]
                    # print("routeValueList", routeValueList)
                # print("cat_*routeValueList",*routeValueList)    
                y = torch.cat(routeValueList,1)
                print("layername_route",layername)
                print("y",y.shape)
                layerValueList.append(y)
                
            # shortcut    
            elif(layertype == "shortcut"):
                # print("shortcut",layersp)
                y = torch.add(layerValueList[layersp[0]],layerValueList[layersp[1]])   
                print("layername_shortcut",layername)
                print("y",y.shape)    
                layerValueList.append(y)
                    # torch.cat()
                # self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps))
            # upsample    
            elif(layertype == "upsample"):
                
                y = self.upsampleLayer(y)
                print("layername_upsample",layername)
                print("y",y.shape)
                layerValueList.append(y)
            # yolohead    
            elif(layertype == "yolohead"):
                y = y.view(-1,self.num_flat_features(y))
                if layersp[0] == 0:
                    y = self.headlinerSmall(y)
                elif layersp[0] == 1:
                    y = self.headlinerMedian(y)
                elif layersp[0] == 2:   
                    y = self.headlinerMax(y)
                y = F.relu(y)
                print("layername_yolohead",layername)
                print("y",y.shape)
                layerValueList.append(y)
                result.append(y)
            # maxpool     
            elif(layertype == "maxpool"):
                # 最大池化，特征图大小不变
                y = F.max_pool2d(y,layersp[0],stride=1, padding=int(layersp[0]/2))
                # y = self.maxPool(y)
                print("layername_maxpool",layername)
                print("y",y.shape)
                layerValueList.append(y)
            
        # x = self.yolov4Layer(x)
        # x = x.view(-1,self.num_flat_features(x))
        # x = F.relu(self.fc1(x))
        return result
        
     #使用num_flat_features函数计算张量x的总特征量（把每个数字都看出是一个特征，即特征总量），比如x是4*2*2的张量，那么它的特征总量就是16。
    def num_flat_features(self, x):
        size = x.size()[1:] # 这里为什么要使用[1:],是因为pytorch只接受批输入，也就是说一次性输入好几张图片，那么输入数据张量的维度自然上升到了4维
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

# yolo卷积块
class Convolutional(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride,padding,momentum,eps):
        super(Convolutional,self).__init__()

        self.convl = nn.Conv2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding)
        self.bn1 = nn.BatchNorm2d(out_channels, momentum = momentum, eps = eps, affine = True)

    def forward(self,x):
        x = self.convl(x)
        x = self.bn1(x) 
        x = F.relu(x)
        return x



# shortcut
# class ShortCut(nn.Module):
#     def __init__(self,in_channels,out_channels,kernel_size,stride,padding,momentum,eps):
#         super(ShortCut,self).__init__()

#     def forward(self,x1,x2):
#         x = torch.add(x1,x2)
#         return x


# route 
# class Route(nn.Module):
#     def __init__(self,in_channels,out_channels,kernel_size,stride,padding,momentum,eps):
#         super(Route,self).__init__()

#     def forward(self,tuple,dim=0):
#         x = torch.cat(**tuple)
#         return x


#  YoloHead 
class YoloHead(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride,padding,momentum,eps):
        super(YoloHead,self).__init__()

        self.convl = nn.Conv2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding)
        self.bn1 = nn.BatchNorm2d(out_channels, momentum = momentum, eps = eps, affine = True)

    def forward(self,x):
        x = self.convl(x)
        x = self.bn1(x) 
        x = F.relu(x)
        return x


 




        